"""
项目：Ai-Engineer自适应工程设计自动化插件
作者：刘昊轩、李佳音、刘伟伟、杨明久（排序不分先后）
©版权保护
"""
# src/ai_integration/deepseek_client_ollama.py
import logging
from typing import Dict, Any
import requests
from dotenv import load_dotenv
from utilities import get_response

load_dotenv()


class OllamaClient:
    """Ollama部署的DeepSeek客户端（主服务）"""

    def __init__(self):
        self.ollama = get_response.Ollama()
        self.model_name = "deepseek-r1:32b"
        self.logger = logging.getLogger("OllamaClient")

        # 验证模型是否存在
        if not self._model_exists():
            raise ModelNotFoundError(f"模型 {self.model_name} 未加载")

    def generate(self, prompt: str, **kwargs) -> str:
        """生成文本"""
        try:
            response = self.ollama.generate(self.model_name, prompt, **kwargs)
            response.raise_for_status()
            return response.json()["response"]
        except requests.exceptions.RequestException as e:
            self.logger.error(f"API请求失败: {str(e)}")
            raise ServiceUnavailableError("Ollama服务不可用") from e

    def _model_exists(self) -> bool:
        """验证模型是否已加载"""
        try:
            response = self.ollama.tags()
            models = [m["name"] for m in response.json()["models"]]
            return self.model_name in models
        except Exception as e:
            raise ConnectionError(f"无法连接Ollama服务: {str(e)}")


class ModelNotFoundError(Exception):
    """模型未找到异常"""
    pass


class ServiceUnavailableError(Exception):
    """服务不可用异常"""
    pass

def get_performance(self) -> Dict[str, Any]:
    """获取服务性能指标"""
    ollama = get_response.Ollama()
    response = ollama.ps()
    return response.json()